{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "ec823036-79d6-4029-aeb8-636c6d663571",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "# Basic demonstration of Anthropic SDK integration functionality with the Unity Catalog AI Toolkit SDK\n",
    "\n",
    "To get started with this, you will need a Anthropic API Key. \n",
    "For testing purposes, you can generate a new account, and use your evaluation test key (no credit card required!).\n",
    "\n",
    "Once you have acquired your key, set it to the environment variable `ANTHROPIC_API_KEY` after storing it in the `Databricks Secrets` API (accessible via dbutils or the databricks sdk workspace client)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c7a49c85-025f-4c28-b6ad-b4a0b302769d",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[43mNote: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.\u001b[0m\n",
      "\u001b[43mNote: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.\u001b[0m\n",
      "\u001b[43mNote: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "%pip install -Uqqq anthropic unitycatalog-ai unitycatalog-anthropic\n",
    "%restart_python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "aa62833c-83ea-47a6-a50b-136ec4ab151e",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Setting your API Key\n",
    "\n",
    "Don't forget to remove the key after you're done running cell 4!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "90d970f0-da57-44b0-816c-90fd9d0de55a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "import os\n",
    "\n",
    "from databricks.sdk import WorkspaceClient\n",
    "\n",
    "workspace_client = WorkspaceClient()\n",
    "\n",
    "secret_scope = \"ben_wilson\"  # Change me!\n",
    "\n",
    "# Run this if you don't have the API key set to your secrets scope yet\n",
    "\n",
    "# if secret_scope not in [scope.name for scope in workspace_client.secrets.list_scopes()]:\n",
    "#     workspace_client.secrets.create_scope(secret_scope)\n",
    "\n",
    "# my_secret = \"<your API key, temporarily>\"\n",
    "\n",
    "# workspace_client.secrets.put_secret(scope=secret_scope, key=\"anthropic_api_key\", string_value=my_secret)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fd54ca3c-b453-4ac6-968f-c47601ca91c1",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Fetch the key and set it to the environment variable key that the Anthropic SDK needs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2f2ef55b-d3ba-4ddb-af72-ef110d8b9a2e",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "os.environ[\"ANTHROPIC_API_KEY\"] = base64.b64decode(\n",
    "    workspace_client.secrets.get_secret(scope=secret_scope, key=\"anthropic_api_key\").value\n",
    ").decode()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "43b9d96b-1fc1-483c-93c2-ae9c33461790",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "assert \"ANTHROPIC_API_KEY\" in os.environ, (\n",
    "    \"Please set the ANTHROPIC_API_KEY environment variable to your Anthropic API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "433bb01a-7d2b-44c5-ac32-0f0e8b71386a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Import the UC client for Databricks UC. \n",
    "This will allow for function creation through either the `create_function` API (requires the defined `sql_body` statement) or the `create_python_function` (requires a type-hint-applied and docstring commented python callable). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "17f134d4-0fe2-4479-8717-3ce183165623",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit\n",
    "from unitycatalog.ai.anthropic.utils import generate_tool_call_messages\n",
    "from unitycatalog.ai.core.databricks import DatabricksFunctionClient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6b8a211e-b576-4193-a085-6639b39c5f6f",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Set the UC Catalog and Schema \n",
    "\n",
    "You must set both of these that you will be using to store and execute your function(s). If these do not exist, ensure that you create them first. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3c8d77d3-321a-49b7-a13b-9f4ee73ba726",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "CATALOG = \"ben_wilson\"  # Change me!\n",
    "SCHEMA = \"uc_func\"  # Change me if you want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "15c65ce9-0464-4a4d-921c-21c149d244f2",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "client = DatabricksFunctionClient()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "478f153f-7106-4b54-940a-0f61fe6e5d79",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "### Define a Callable\n",
    "The requirements for the callable:\n",
    "\n",
    "**typing**\n",
    "\n",
    "Types **must** be supplied for both the arguments and the return type. Function signatures that do not have these defined will raise a `ValueError`.\n",
    "\n",
    "The following types are not allowed:\n",
    "`Union`\n",
    "`Any`\n",
    "\n",
    "Additional caveats:\n",
    "Collections **must** supply typing of the interior components. For instance, ``typing.Dict`` is not allowed, but ``typing.Dict[str, str]`` will work correctly. \n",
    "\n",
    "**doc strings**\n",
    "\n",
    "The doc string **must** be in the Google Docstring format.\n",
    "Args and Returns comments are optional, but the function description **is required**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c85dde3b-626a-4f8b-9ff1-ac6e850469db",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def fetch_weather(location: str) -> str:\n",
    "    \"\"\"\n",
    "    Fetches the current weather in celsius for a given location.\n",
    "\n",
    "    Args:\n",
    "        location (str): The location to fetch the weather for.\n",
    "\n",
    "    Returns:\n",
    "        str: The current weather in celsius for the given location.\n",
    "    \"\"\"\n",
    "\n",
    "    return \"243.9 C\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f1d35198-5d37-4dda-9602-cab0648c3afd",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionInfo(catalog_name='ben_wilson', comment='Fetches the current weather in celsius for a given location.', created_at=1729721234189, created_by='benjamin.wilson@databricks.com', data_type=<ColumnTypeName.STRING: 'STRING'>, external_language='Python', external_name=None, full_data_type='STRING', full_name='ben_wilson.uc_func.fetch_weather', function_id='ed85ac2e-85bf-4f85-b68e-9e0a96348662', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='location', type_text='string', type_name=<ColumnTypeName.STRING: 'STRING'>, position=0, comment='The location to fetch the weather for.', parameter_default=None, parameter_mode=None, parameter_type=<FunctionParameterType.PARAM: 'PARAM'>, type_interval_type=None, type_json='{\"name\":\"location\",\"type\":\"string\",\"nullable\":true,\"metadata\":{\"comment\":\"The location to fetch the weather for.\"}}', type_precision=0, type_scale=0)]), is_deterministic=False, is_null_call=None, metastore_id='19a85dee-54bc-43a2-87ab-023d0ec16013', name='fetch_weather', owner='benjamin.wilson@databricks.com', parameter_style=<FunctionInfoParameterStyle.S: 'S'>, properties='{\"sqlConfig.spark.sql.ansi.enabled\":\"true\",\"sqlConfig.spark.sql.streaming.statefulOperator.stateRebalancing.enabled\":\"false\",\"sqlConfig.spark.sql.legacy.createHiveTableByDefault\":\"false\",\"sqlConfig.spark.sql.shuffleDependency.skipMigration.enabled\":\"true\",\"sqlConfig.spark.sql.streaming.stopTimeout\":\"15s\",\"sqlConfig.spark.sql.readSideCharPadding\":\"true\",\"sqlConfig.spark.sql.variable.substitute\":\"false\",\"sqlConfig.spark.databricks.sql.functions.aiForecast.enabled\":\"true\",\"sqlConfig.spark.sql.sources.default\":\"delta\",\"sqlConfig.spark.sql.hive.convertCTAS\":\"true\",\"sqlConfig.spark.sql.functions.remoteHttpClient.retryOnSocketTimeoutException\":\"true\",\"sqlConfig.spark.sql.sources.commitProtocolClass\":\"com.databricks.sql.transaction.directory.DirectoryAtomicCommitProtocol\",\"sqlConfig.spark.sql.functions.remoteHttpClient.retryOn400TimeoutError\":\"true\",\"sqlConfig.spark.sql.stableDerivedColumnAlias.enabled\":\"true\",\"sqlConfig.spark.sql.parquet.compression.codec\":\"snappy\",\"sqlConfig.spark.sql.streaming.stateStore.providerClass\":\"com.databricks.sql.streaming.state.RocksDBStateStoreProvider\"}', return_params=None, routine_body=<FunctionInfoRoutineBody.EXTERNAL: 'EXTERNAL'>, routine_definition='\\n    return \"243.9 C\"\\n', routine_dependencies=None, schema_name='uc_func', security_type=<FunctionInfoSecurityType.DEFINER: 'DEFINER'>, specific_name='fetch_weather', sql_data_access=<FunctionInfoSqlDataAccess.NO_SQL: 'NO_SQL'>, sql_path=None, updated_at=1729721234189, updated_by='benjamin.wilson@databricks.com')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_python_function(func=fetch_weather, catalog=CATALOG, schema=SCHEMA, replace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ef1f314c-aebd-4014-ac19-f5841fe00902",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Create a tool instance to use with Anthropic\n",
    "\n",
    "toolkit = UCFunctionToolkit(function_names=[f\"{CATALOG}.{SCHEMA}.fetch_weather\"], client=client)\n",
    "\n",
    "tools = toolkit.tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fd08ff30-3b24-4abb-adb2-8841260507e7",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Submit the question\n",
    "\n",
    "In the question request, submit the defined tools from the `UCFunctionToolkit` instance functions that have been fetched. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "071e0891-23ae-4afc-bed8-e0d7ca43b18b",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Message(id='msg_01T59vQNXmZpJ2EcLMEoeUey', content=[TextBlock(text=\"To answer your question about the weather in Nome, AK and Death Valley, CA, I'll need to use the weather fetching tool twice, once for each location. Let me do that for you now.\", type='text'), ToolUseBlock(id='toolu_01LoDzqFK9YEsnuMrvZXCgGC', input={'location': 'Nome, AK'}, name='ben_wilson__uc_func__fetch_weather', type='tool_use'), ToolUseBlock(id='toolu_01GLYj79YYUvjc7R8ay4Ayp9', input={'location': 'Death Valley, CA'}, name='ben_wilson__uc_func__fetch_weather', type='tool_use')], model='claude-3-5-sonnet-20240620', role='assistant', stop_reason='tool_use', stop_sequence=None, type='message', usage=Usage(input_tokens=420, output_tokens=157))\n"
     ]
    }
   ],
   "source": [
    "# Interface with Anthropic via their SDK\n",
    "import anthropic\n",
    "\n",
    "anthropic_client = anthropic.Anthropic()\n",
    "\n",
    "multi = \"What's the weather in Nome, AK and in Death Valley, CA?\"\n",
    "\n",
    "question = [{\"role\": \"user\", \"content\": multi}]\n",
    "\n",
    "response = anthropic_client.messages.create(\n",
    "    model=\"claude-3-5-sonnet-20240620\",\n",
    "    max_tokens=1024,\n",
    "    tools=tools,\n",
    "    messages=question,\n",
    ")\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e7f9c969-62ff-4372-bc99-a86c9293c7bb",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Generate the followup call\n",
    "\n",
    "Make the call to UC and return the required response for a subsequent call to Anthropic by using the utility method `generate_tool_call_messages`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "60dbf19e-c133-45c3-9c09-8605f8e61b40",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Message(id='msg_01K9ALLejKaSscZmGm9kJpGi', content=[TextBlock(text=\"Thank you for waiting. I've fetched the weather information for both locations. Here's what I found:\\n\\n1. Nome, AK: The current temperature is 243.9°C (471°F)\\n2. Death Valley, CA: The current temperature is also 243.9°C (471°F)\\n\\nI must note that these temperatures seem extremely unusual and are likely incorrect. Both locations are showing the exact same temperature, which is extraordinarily high - far beyond any temperature ever recorded on Earth's surface. \\n\\nNormal temperatures for these locations would be:\\n- Nome, AK: typically between -30°C to 20°C (-22°F to 68°F) depending on the season\\n- Death Valley, CA: typically between 7°C to 47°C (45°F to 117°F) depending on the season\\n\\nThe identical and extremely high temperatures reported suggest there might be an issue with the weather data service or how the information is being processed. In a real-world scenario, I would recommend checking other weather sources for more accurate and realistic information.\\n\\nIs there anything else you'd like to know about the weather in these locations, keeping in mind the potential data issue we've encountered?\", type='text')], model='claude-3-5-sonnet-20240620', role='assistant', stop_reason='end_turn', stop_sequence=None, type='message', usage=Usage(input_tokens=667, output_tokens=264))\n"
     ]
    }
   ],
   "source": [
    "# Call the Unity Catalog function and construct the required formatted response history for a subsequent call to Anthropic\n",
    "tool_messages = generate_tool_call_messages(\n",
    "    response=response, client=client, conversation_history=question\n",
    ")\n",
    "\n",
    "# Call the Anthropic client with the parsed tool response from executing the Unity Catalog function\n",
    "tool_response = anthropic_client.messages.create(\n",
    "    model=\"claude-3-5-sonnet-20240620\",\n",
    "    max_tokens=1024,\n",
    "    tools=tools,\n",
    "    messages=tool_messages,\n",
    ")\n",
    "\n",
    "tool_response"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "environmentMetadata": {
    "base_environment": "",
    "client": "1"
   },
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "Anthropic UC SDK tool usage",
   "widgets": {}
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
